{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Huấn luyện có giám sát với SFTTrainer\n",
    "\n",
    "Bài học này se dạy bạn các huấn luyện mô hình `HuggingFaceTB/SmolLM2-135M` bằng `SFTTrainer` trong thư viện `trl`.  Các cell trong notebook này sẽ chạy và huấn luyện mô hình. Bạn có thể chọn độ khó bằng cách thử nghiệm với các bộ dữ liệu khác nhau.\n",
    "\n",
    "<div style='background-color: lightblue; padding: 10px; border-radius: 5px; margin-bottom: 20px; color:black'>\n",
    "    <h2 style='margin: 0;color:blue'>Bài tập: Fine-Tuning SmolLM2 với SFTTrainer</h2>\n",
    "    <p>Chọn một bộ dự liệu từ Hugging Face hub và huấn luyện một mô hình trên bộ dữ liệu đó. </p> \n",
    "    <p><b>Các bài tập</b></p>\n",
    "    <p>🐢 Sử dụng bộ dữ liệu `HuggingFaceTB/smoltalk`</p>\n",
    "    <p>🐕 Thử nghiệm với bộ dữ liệu `bigcode/the-stack-smol` và huấn luyện một mô hình sinh code trên tập con cụ thể `data/python`.</p>\n",
    "    <p>🦁 Chọn một bộ dữ liệu liên quan đến một lĩnh vực mà bạn quan tâm</p>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Cài đặt các thư viện cần thiết\n",
    "# !pip install transformers datasets trl huggingface_hub\n",
    "\n",
    "# Đăng nhập vào Hugging Face\n",
    "from huggingface_hub import login\n",
    "\n",
    "login()\n",
    "\n",
    "# Để thuận tiện, bạn có thể tạo một biến môi trường chứa `token hub` của bạn dưới dạng HF_TOKEN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Các thư viện cần thiết\n",
    "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
    "from datasets import load_dataset\n",
    "from trl import SFTConfig, SFTTrainer, setup_chat_format\n",
    "import torch\n",
    "\n",
    "device = (\n",
    "    \"cuda\"\n",
    "    if torch.cuda.is_available()\n",
    "    else \"mps\" if torch.backends.mps.is_available() else \"cpu\"\n",
    ")\n",
    "\n",
    "# Tải mô hình và tokenizer\n",
    "model_name = \"HuggingFaceTB/SmolLM2-135M\"\n",
    "model = AutoModelForCausalLM.from_pretrained(\n",
    "    pretrained_model_name_or_path=model_name\n",
    ").to(device)\n",
    "tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path=model_name)\n",
    "\n",
    "# Thiết lập định dạng chat\n",
    "model, tokenizer = setup_chat_format(model=model, tokenizer=tokenizer)\n",
    "\n",
    "# Đặt tên cho mô hình huấn luyện để lưu &/ tải lên\n",
    "finetune_name = \"SmolLM2-FT-MyDataset\"\n",
    "finetune_tags = [\"smol-course\", \"module_1\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Sinh văn bản với Mô hình gốc\n",
    "\n",
    "Ở đây chúng ta sẽ thử nghiệm mô hình gốc chưa được huấn luyện trên định dạng chat."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Kiểm tra mô hình gốc trước khi huấn luyện\n",
    "prompt = \"Write a haiku about programming\"\n",
    "\n",
    "# Định dạng\n",
    "messages = [{\"role\": \"user\", \"content\": prompt}]\n",
    "formatted_prompt = tokenizer.apply_chat_template(messages, tokenize=False)\n",
    "\n",
    "# Tạo phản hồi từ mô hình\n",
    "inputs = tokenizer(formatted_prompt, return_tensors=\"pt\").to(device)\n",
    "outputs = model.generate(**inputs, max_new_tokens=100)\n",
    "print(\"Before training:\")\n",
    "print(tokenizer.decode(outputs[0], skip_special_tokens=True))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Chuẩn bị dữ liệu\n",
    "\n",
    "Chúng ta sẽ tải một bộ dữ liệu mẫu và định dạng nó cho việc huấn luyện. Bộ dữ liệu cần được cấu trúc với các cặp đầu vào - đầu ra, trong đó mỗi đầu vào là một chỉ thị và đầu ra là phản hồi mong đợi từ mô hình.\n",
    "\n",
    "**TRL sẽ định dạng các tin nhắn đầu vào dựa trên định dạng chat của mô hình** Chúng cần được biểu diễn dưới dạng danh sách các từ điển với các khóa: `role` và `content`.\n",
    "\n",
    "**Ví dụ:**\n",
    "```sh\n",
    "[\n",
    "  {\"role\": \"user\", \"content\": \"Hello, how are you?\"},\n",
    "  {\"role\": \"assistant\", \"content\": \"I'm doing well, thank you! How can I assist you today?\",},\n",
    "]\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Tải dữ liệu mẫu\n",
    "from datasets import load_dataset\n",
    "\n",
    "# TODO: Tải bộ dữ liệu thông qua việc điều chỉnh các tham số path và name\n",
    "ds = load_dataset(path=\"HuggingFaceTB/smoltalk\", name=\"everyday-conversations\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: 🦁 Nếu dataset của bạn không ở định dạng mà TRL có thể chuyển đổi thành định dạng chat, bạn sẽ cần xử lý nó.\n",
    "# Tham khảo [Định dạng Chat](../chat_templates.md)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Điều chỉnh SFTTrainer\n",
    "\n",
    "Điều chỉnh `SFTTrainer` với các tham số khác nhau giúp điều khiển quá trình huấn luyện trở nên hiệu quả hơn. Các thông số bao gồm\n",
    "- Số bước huấn luyện (steps)\n",
    "- Kích thước batch (batch size)\n",
    "- Tốc độ học (learning rate)\n",
    "- Chiến lược đánh giá mô hình (evaluation strategy)\n",
    "\n",
    "Ngoài ra, còn rất nhiều thông số khác, bạn có thể tham khảo thêm ở [SFTTrainer](https://huggingface.co/docs/trl/sft_trainer)\n",
    "\n",
    "Điều chỉnh các tham số này dựa trên yêu cầu cụ thể và tài nguyên tính toán của bạn."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Điều chỉnh SFTTrainer\n",
    "sft_config = SFTConfig(\n",
    "    output_dir=\"./sft_output\",\n",
    "    max_steps=1000,  # Điều chỉnh dựa trên kích thước dataset và thời lượng huấn luyện mong muốn\n",
    "    per_device_train_batch_size=4,  # Đặt theo dung lượng bộ nhớ GPU của bạn\n",
    "    learning_rate=5e-5,  # Phổ biến cho quá trình huấn luyện có giám sát\n",
    "    logging_steps=10,  # Tần suất ghi log các metrics huấn luyện\n",
    "    save_steps=100,  # Tần suất lưu các checkpoint mô hình\n",
    "    evaluation_strategy=\"steps\",  # Đánh giá mô hình theo các khoảng thời gian\n",
    "    eval_steps=50,  # Tần suất đánh giá\n",
    "    use_mps_device=(\n",
    "        True if device == \"mps\" else False\n",
    "    ),  # Sử dụng MPS cho huấn luyện độ chính xác hỗn hợp\n",
    "    hub_model_id=finetune_name,  # Đặt tên cho mô hình của bạn\n",
    ")\n",
    "\n",
    "# Khởi tạo SFTTrainer\n",
    "trainer = SFTTrainer(\n",
    "    model=model,\n",
    "    args=sft_config,\n",
    "    train_dataset=ds[\"train\"],\n",
    "    tokenizer=tokenizer,\n",
    "    eval_dataset=ds[\"test\"],\n",
    ")\n",
    "\n",
    "\n",
    "# TODO: 🦁 🐕 căn chỉnh các tham số SFTTrainer với bộ dữ liệu bạn đã chọn.\n",
    "# Ví dụ, nếu bạn đang sử dụng bộ `bigcode/the-stack-smol`, bạn sẽ cần chọn cột `content`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Huấn luyện Mô hình\n",
    "\n",
    "Với trainer đã được điều chỉnh, chúng ta có thể tiến hành huấn luyện mô hình. Quá trình huấn luyện sẽ bao gồm\n",
    "- Lặp qua bộ dữ liệu\n",
    "- Tính toán loss\n",
    "- Cập nhật các tham số của mô hình để giảm thiểu loss này."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Huấn luyện mô hình\n",
    "trainer.train()\n",
    "\n",
    "# Lưu mô hình\n",
    "trainer.save_model(f\"./{finetune_name}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Đưa mô hình lên Hugging Face Hub\n",
    "trainer.push_to_hub(tags=finetune_tags)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div style='background-color: lightblue; padding: 10px; border-radius: 5px; margin-bottom: 20px; color:black'>\n",
    "    <h2 style='margin: 0;color:blue'>Bài tập thêm: Sinh văn bản với mô hình vừa được huấn luyện</h2>\n",
    "    <p>🐕 Sử dụng mô hình đã được huấn luyện để sinh ra phản hồi, giống như với ví dụ ban đầu.</p>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Kiểm tra mô hình đã được huấn luyện trên cùng một yêu cầu\n",
    "\n",
    "# Kiểm tra mô hình gốc trước khi huấn luyện\n",
    "prompt = \"Write a haiku about programming\"\n",
    "\n",
    "# Định dạng chat\n",
    "messages = [{\"role\": \"user\", \"content\": prompt}]\n",
    "formatted_prompt = tokenizer.apply_chat_template(messages, tokenize=False)\n",
    "\n",
    "# Sinh phản hồi từ mô hình\n",
    "inputs = tokenizer(formatted_prompt, return_tensors=\"pt\").to(device)\n",
    "\n",
    "# TODO: sử dụng mô hình đã được huấn luyện để sinh phản hồi, giống như với ví dụ."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 💐 Chúc mừng bạn. Bạn đã hoàn thành!\n",
    "\n",
    "Bài tập này đã cung cấp hướng dẫn từng bước để bạn huấn luyện được mô hình `HuggingFaceTB/SmolLM2-135M` sử dụng `SFTTrainer`. Bằng cách làm theo các bước này, bạn có thể điều chỉnh mô hình để thực hiện các tác vụ cụ thể hiệu quả hơn. Nếu bạn muốn tiếp tục làm việc với khóa học này, đây là một số bước bạn có thể thử:\n",
    "\n",
    "- Thử notebook này ở mức độ khó hơn\n",
    "- Review PR của học viên khác\n",
    "- Cải thiện tài liệu khóa học thông qua Issue hoặc PR."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "py310",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
